# 将静态手势打包，仿照cifar.npz的形式
import os
import numpy as np
from PIL import Image


# 静态数字手势的图片文件夹，下载后放在此子项目文件夹之下，即"PrepareGestures\gesture_static_number"
# 数据集有一个问题要注意，直接读取原始数据集后在线缩小到32*32，生成的npz文件在训练时会loss nan，而将原图像在额外程序中线下缩小到32*32，则可训练
# 因此建议使用其他手段线下缩小原文件生成新文件，此处则直接读取缩放好的图像文件
dataset_dir = 'gesture_static_number'


# 根据图片文件列表读取图片内容，并根据文件名建立标签，返回numpy数组
def imgs_to_numpy(dirs):
	imgs = np.array([], dtype = np.uint8)
	labels = []
	for dir in dirs:
		img = Image.open('gesture_static_number/' + dir)
		# img = img.resize((32, 32), Image.BILINEAR) # 直接读取尺寸为32*32的文件，不要读取原始224*171的文件并在线缩小
		img_arr = np.asarray(img)
		img_arr = np.reshape(img_arr, [-1, 32*32])
		if (imgs.size == 0):
			imgs = img_arr
		else:
			imgs = np.vstack((imgs, img_arr))

		# 注意label值从0开始
		labels.append(int(dir[0:2]) - 1)

	return imgs, np.array(labels, dtype = np.int8)


# 读取图片文件列表，返回一个打乱顺序的训练集图片文件名列表和一个打乱顺序的验证集图片文件名列表，二者数量比4:1
def random_imgs_list():
	train_list = []
	valid_list = []
	dirs = os.listdir(dataset_dir)
	img_nums = len(dirs)
	for i in range(0,5):
		train_list += dirs[i * 400 : i * 400 + 320]
		valid_list += dirs[i * 400 + 320 : i * 400 + 400]
	np.random.shuffle(train_list)
	np.random.shuffle(valid_list)
	return train_list, valid_list


def main():
	train_dirs, valid_dirs = random_imgs_list()
	train_images, train_labels = imgs_to_numpy(train_dirs)
	validation_images, validation_labels = imgs_to_numpy(valid_dirs)
	np.savez_compressed('gesture', train_images=train_images, validation_images=validation_images,
                    train_labels=train_labels, validation_labels=validation_labels)


if __name__ == '__main__':
    main()
